import asyncio
import json
from pathlib import Path
from typing import Any, Dict

from langchain.agents import initialize_agent, AgentType
from langchain_mcp_adapters.client import MultiServerMCPClient

from a0base.base_llm import pop_llm

# ---------- 工具函数 ----------
CONFIG_FILE = Path("weather_mcp_config.json")


def load_servers(file_path: str | Path = CONFIG_FILE) -> Dict[str, Any]:
    """
    读取 MCP 服务器配置
    :param file_path: 配置文件路径
    :return: mcpServers 字段内容（dict）
    """
    with open(file_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    return data.get("mcpServers", {})


async def run_chat_loop() -> None:
    """
    启动 DeepSeek + MCP 的命令行聊天循环
    """
    # 1️⃣ 加载配置
    servers_cfg = load_servers()

    # 2️⃣ 初始化 MCP 客户端并拿工具
    mcp_client = MultiServerMCPClient(servers_cfg)
    tools = await mcp_client.get_tools()
    print(f"[INFO] 已加载 {len(tools)} 个 MCP 工具：{[t.name for t in tools]}")

    # 3️⃣ 构建 Agent
    agent = initialize_agent(
        tools=tools,
        llm=pop_llm,
        agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
    )

    # 4️⃣ CLI 循环
    print("\n🤖 DeepSeek-MCP 聊天启动，输入 'quit' 退出。\n")
    while True:
        user_input = input("你: ").strip()
        if user_input.lower() == "quit":
            break
        try:
            answer = await agent.arun(user_input)
            print(f"AI: {answer}")
        except Exception as exc:
            print(f"[ERROR] 调用失败: {exc}")

    print("👋 会话已结束，Bye!")


# ---------- 入口 ----------
if __name__ == "__main__":
    asyncio.run(run_chat_loop())
